#!/usr/bin/env python

import os, sys, logging
import hashlib, datetime

logfname = '/tmp/dmq_ioloop_test.log'
loglevel = logging.DEBUG
loglevel = logging.INFO
logging.basicConfig(filename=logfname,
			format='%(asctime)s %(levelname)s %(message)s',
			level=loglevel)
#logging.basicConfig(filename='/tmp/dmq_ioloop_test.log', 
#			format='%(asctime)s %(levelname)s %(message)s',
#			level=logging.INFO)

SITE_ROOT = os.path.dirname(os.path.realpath(__file__))
path = os.path.normpath(SITE_ROOT + '/../')
if path not in sys.path:
	sys.path.append(path)

import unittest, time
import zmq, dmq.ioloop
from dmq import znode

class TestIOLoop(unittest.TestCase):

	def _start_ioloop(self):
		logging.debug("_start_ioloop")
		"""docstring for _start_ioloop"""

		self.ioloop.start()
	#_start_ioloop()

	def set_up_push_sync(self, send_func=None):
		"""docstring for set_up_push_sync"""
		self.sync_push_znode = znode.ZNodePUSH(port=self.sync_port).add()
		self.sync_push_znode.on_send(send_func)
	#set_up_push_sync()

	def set_up_pull_sync(self, app=None, recv_func=None):
		"""docstring for set_up_pull_sync"""

		def pull_sync_app(env, msg):
			"""docstring for pull_sync_app"""
			self.test_msg = msg
		#pull_sync_app()

		if app:
			self.sync_pull_znode = znode.ZNodePULL(app=app,
					port=self.sync_port).add()
		else:
			self.sync_pull_znode = znode.ZNodePULL(app=set_up_pull_sync,
					port=self.sync_port).add()

		self.sync_pull_znode.on_recv(recv_func)
	#set_up_pull_sync()

	def setUp(self):
		"""
		Start and fork a new process that will contain an ioloop instance for
		our tests. The new process will run untill it is killed with a SIGTERM
		in tearDown().
		"""
		logging.debug("setUp()")

		self.zmq_host = '*'
		self.zmq_transport = 'tcp'
		self.zmq_context = None
		self.zmq_sockets = []
		self.ioloop = None
		self.ioloop = dmq.ioloop.IOLoop.instance()
		self._child_pid = 0
		self.opcount = 100
		self.op_gen = range(self.opcount)
		self.sync_port = 9666
		self.sync_push_znode = None
		self.sync_pull_znode = None
		self.test_msg = "A test message" 

		logging.debug("setUp() get a context")
		self.zmq_context = zmq.Context()
	#setUp()

	def _gen_msg(self, token=None):
		"""docstring for _gen_msg"""
		#logging.debug("TestIOLoop._gen_msg(token:%s) op_gen:%s" % (token, self.op_gen))

		if not token:
			try:
				token = self.op_gen.pop(0)
			except IndexError as e:
				return None

		try:
			self.assertIsNotNone(token)
		except:
			self.ioloop.stop()
			raise
		msg = "%s:%s" % (self.test_msg, token)

		return "%s:%s" % (hashlib.sha1(msg).hexdigest(), msg)
	#_gen_msg()

	def tearDown(self):
		logging.debug("tearDown")
		"""docstring for tearDown"""
		self.ioloop.stop()
	#tearDown()

	def runTest(self):
		"""docstring for runTest"""
		# Verify the sync nodes exist
		self.assertIsInstance(znode.ZNodePUSH, self.sync_push_znode)
		self.assertIsInstance(znode.ZNodePULL, self.sync_pull_znode)
	#runTest()

#TestIOLoop

class TestIOLoopReqRepEcho(TestIOLoop):

	def setUp(self):
		self.reqrep_port = 5678
		self.req_node = None
		self.rep_node = None

		super(TestIOLoopReqRepEcho, self).setUp()
	#setUp()

	def test_req(self):
		"""docstring for req_test"""
		logging.debug("TestIOLoopReqRepEcho.test_req()")

		expected_msg = self._gen_msg()
		def bootstrap(env):
			"""docstring for bootstrap"""
			logging.debug("TestIOLoopReqRepEcho.test_req.bootstrap()")
			# Get things started.
			self.req_node.send(expected_msg)
			return expected_msg
		#bootstrap()

		def send_stop(env):
			"""docstring for send_stop"""
			self.ioloop.stop()
			return 'stop'
		#send_stop()

		def req_app(environ, msg):
			logging.debug("TestIOLoopReqRepEcho.test_req.req_app() msg : %s" % (msg,))

			self.assertEqual(req_app.expected_msg, msg)

			req_app.expected_msg = self._gen_msg()
			if not req_app.expected_msg:
				logging.debug("TestIOLoopReqRepEcho.test_req.req_app() start stopping the test")
				self.sync_push_znode.on_send(send_stop)
				return None

			self.sync_push_znode.send(req_app.expected_msg)

			return req_app.expected_msg
		#req_app()
		req_app.expected_msg = expected_msg

		self.req_node = znode.znode(req_app, type=zmq.REQ, port=self.reqrep_port)
		self.set_up_push_sync(bootstrap)
		self._start_ioloop()
	#test_req()

	def test_rep(self):
		"""docstring for rep_test"""

		def pull_app(env, msg):
			"""use the pull to syncronize our req/rep messages and test them"""
			logging.debug("TestIOLoopReqRepEcho.rep_test.pull_app() msg : %s" % (msg,))

			if msg == 'stop':
				logging.debug("TestIOLoopReqRepEcho.rep_test.pull_app() stopping ioloop")
				self.ioloop.stop()
				return
				
			self.test_msg = msg
			self.rep_node.on_recv(rep_app)
		#pull_app()

		def rep_app(environ, msg):
			logging.debug("TestIOLoopReqRepEcho.rep_test.rep_app() msg : %s" % (msg,))

			self.assertEqual(msg, self.test_msg)
			return self.test_msg
		#rep_app()

		logging.debug("TestIOLoopReqRepEcho.rep_test() ")
		# This is a responder for the req test, so use the req port
		self.set_up_pull_sync(app=pull_app)
		self.rep_node = znode.znode(app=None, type=zmq.REP, port=self.reqrep_port)

		logging.debug("TestIOLoopReqRepEcho.rep_test() starting ioloop")
		self._start_ioloop()
	#test_rep()

#TestIOLoopReqRepEcho

class TestIOLoopPushPull(TestIOLoop):

	def setUp(self):
		self.pushpull_port = 5674
		self.push_node = None
		self.pull_node = None

		super(TestIOLoopPushPull, self).setUp()
	#setUp()

	def test_push(self):
		"""docstring for test_push"""

		def pull_msg(env, msg):
			"""docstring for pull_msg"""
			logging.debug("TestIOLoopPushPull.test_push.pull_msg() recvd : %s" % (msg,))

			if msg == 'stop':
				logging.debug("TestIOLoopPushPull.test_push.pull_msg() start stopping the test")
				self.ioloop.stop()
				return
				
			expected_msg = self._gen_msg()
			self.assertEqual(expected_msg, msg)

			logging.debug("TestIOLoopPushPull.test_push.pull_msg() sending %s" % (expected_msg,))
			self.push_node.send(expected_msg)
		#pull_msg()

		logging.debug("TestIOLoopPushPull.test_push() starting the test")
		self.push_node = znode.znode(type=zmq.PUSH, port=self.pushpull_port)
		self.set_up_pull_sync(pull_msg)
		self._start_ioloop()
	#test_push()

	def test_pull(self):
		"""docstring for test_pull"""

		def bootstrap(env):
			"""docstring for bootstrap"""
			pull_app.msg = self._gen_msg()

			logging.debug("TestIOLoopPushPull.test_pull.bootstrap() bootstrap msg %s" % (pull_app.msg,))
			return pull_app.msg
		#bootstrap()

		def pull_app(env, msg):
			"""docstring for pull_appq"""
			logging.debug("TestIOLoopPushPull.test_pull.pull_app() msg %s" % (msg,))

			self.assertEqual(msg, pull_app.msg)
			pull_app.msg = self._gen_msg()
			if not pull_app.msg:
				def stop_test(env):
					"""docstring for stop_test"""
					logging.debug("TestIOLoopPushPull.test_pull.pull_app.stop_test() stopping test")
					self.ioloop.stop()
					return 'stop'
				#stop_test()
				logging.debug("TestIOLoopPushPull.test_pull.pull_app() out of test msgs")
				self.sync_push_znode.on_send(stop_test)
				return
				
			logging.debug("TestIOLoopPushPull.test_pull.pull_app() sending msg %s" % (pull_app.msg,))
			self.sync_push_znode.send(pull_app.msg)
		#pull_app()

		self.pull_node = znode.znode(app=pull_app, type=zmq.PULL, port=self.pushpull_port)
		self.set_up_push_sync(bootstrap)
		self._start_ioloop()
	#test_pull()

#TestIOLoopPushPull

class TestIOLoopPubSub(TestIOLoop):

	def setUp(self):
		self.pubsub_port = 5676
		self.pub_node = None
		self.sub_node = None
		super(TestIOLoopPubSub, self).setUp()
	#setUp()

	def test_sub(self):
		"""docstring for test_sub"""

		def bootstrap():
			"""docstring for bootstrap"""
			sub_app.msg = self._gen_msg()
			logging.debug("TestIOLoopPubSub.test_sub.bootstrap() bootstrap msg %s" % (sub_app.msg,))
			self.sub_node.subscribe("")
			self.sub_node.update(znode.ZNode.POLLIN,rem=True)
			self.sub_node.update(znode.ZNode.POLLIN)
			self.sync_push_znode.send(sub_app.msg)
			#return sub_app.msg
		#bootstrap()

		def sub_app(env, msg):
			"""docstring for sub_app"""
			logging.debug("TestIOLoopPubSub.test_sub.sub_app() msg %s" % (msg,))

			self.assertEqual(msg, sub_app.msg)
			sub_app.msg = self._gen_msg()
			if not sub_app.msg:
				def stop_test(env):
					"""docstring for stop_test"""
					logging.debug("TestIOLoopPubSub.test_sub.sub_app.stop_test() stopping test")
					self.ioloop.stop()
					return 'stop'
				#stop_test()
				logging.debug("TestIOLoopPubSub.sub_app.pull_app() out of test msgs")
				self.sync_push_znode.on_send(stop_test)
				return
				
			logging.debug("TestIOLoopPubSub.test_sub.sub_app() sending msg %s" % (sub_app.msg,))
			self.sync_push_znode.send(sub_app.msg)
		#sub_app()

		logging.debug("TestIOLoopPubSub.test_sub.() starting ioloop")
		#self.set_up_push_sync(bootstrap)
		self.ioloop.add_timeout(datetime.timedelta(0,1), bootstrap)
		self.set_up_push_sync()
		self.sub_node = znode.znode(app=sub_app, type=zmq.SUB, port=self.pubsub_port)
		self._start_ioloop()
	#test_sub()

	def test_pub(self):
		"""docstring for test_pub"""

		def pull_msg(env, msg):
			"""docstring for pull_msg"""
			logging.debug("TestIOLoopPubSub.test_pub.pull_msg() recvd : %s" % (msg,))

			#if not self.pub_node:
			#	logging.debug("TestIOLoopPubSub.test_pub() creating pub node.")
			#	self.pub_node = znode.znode(type=zmq.PUB, port=self.pubsub_port)
				
			if msg == 'stop':
				logging.debug("TestIOLoopPubSub.test_pub.pull_msg() start stopping the test")
				self.ioloop.stop()
				return
				
			expected_msg = self._gen_msg()
			self.assertEqual(expected_msg, msg)

			logging.debug("TestIOLoopPubSub.test_pub.pull_msg() sending %s" % (expected_msg,))
			self.pub_node.send(expected_msg)
		#pull_msg()

		def bootstrap():
			"""docstring for bootstrap"""
			self.set_up_pull_sync(pull_msg)
		#bootstrap()

		logging.debug("TestIOLoopPubSub.test_pub() starting the test")

		# add a timeout to allow for subscriber to connect
		#self.ioloop.add_timeout(datetime.timedelta(0,1), bootstrap)
		self.pub_node = znode.znode(type=zmq.PUB, port=self.pubsub_port)
		self.set_up_pull_sync(pull_msg)
		self._start_ioloop()
	#test_pub()
#TestIOLoopPubSub

#class TestIOLoopDealerNode(TestIOLoop):
#
#	def _start_ioloop(self):
#		super(TestIOLoopRequestNode, self)._start_ioloop()
#	#_start_ioloop()
#
#	def test_dealer(self):
#		"""docstring for test_dealer"""
#		pass
#	#test_dealer()
##TestIOLoopDealerNode
#
#class TestIOLoopRouterNode(TestIOLoop):
#
#	def _start_ioloop(self):
#		super(TestIOLoopRequestNode, self)._start_ioloop()
#	#_start_ioloop()
#
#	def test_router(self):
#		"""docstring for test_router"""
#		pass
#	#test_router()
##TestIOLoopRouterNode

if __name__ == '__main__':
	unittest.main()
